import libsvm.intex_node;
import libsvm.svm;
import libsvm.svm_model;
import libsvm.svm_parameter;
import libsvm.svm_problem;

import java.util.*;
import java.io.*;

public class IntexModel {

  svm_parameter m_params;
  svm_problem m_problem;
  svm_model m_model;

  static public int MAX_LENGTH = 10;

  public IntexModel()
  {
  }


  public boolean trainModel(String corpus_dir, String corpus_list)
  {
    readProblem(corpus_dir, corpus_list);
    setParams();

    String strError = svm.svm_check_parameter(m_problem, m_params);
    if(strError != null) {
      System.err.println("Error: " + strError +"!");
      return false;
    }

    m_model = svm.svm_train(m_problem, m_params);

    return true;
  }


  protected boolean readProblem(String corpus_dir, String corpus_list)
  {
    Vector vx = new Vector();
    Vector vy = new Vector();
    int[] counts = new int[2];

    try {
      BufferedReader br = new BufferedReader(new FileReader(corpus_list));
      String strLine = br.readLine();
      while (strLine != null) {
        strLine = strLine.trim();
        if (strLine.length() > 0) {
          // load input document
          String strInFile = corpus_dir + "/" + strLine;
          
          System.out.println("processing [" + strInFile + "] ... ");
	  
          ACEDocument aceDoc = new ACEDocument();
          aceDoc.load(strInFile);
	  aceDoc.removeRepetitions();
          ACETokSentDocument aceTSDoc = new ACETokSentDocument(aceDoc);
          
          // create relation extraction document
          IntexDocument intexDoc = new IntexDocument(aceTSDoc);
	  // store all relation examples in vx, vy
	  updateProblem(intexDoc, vx, vy);
	  // update counts
	  int[] dCounts = intexDoc.pairsCounts();
	  for (int i = 0; i < counts.length; i++)
	    counts[i] += dCounts[i];

	  //	  System.out.println("done!");
        }	  
        strLine = br.readLine();
      }
    }
    catch (IOException e) {
      System.err.println(e);
      return false;
    }

    m_problem = new svm_problem();
    m_problem.l = vy.size();
    m_problem.x = (intex_node[][]) vx.toArray(new intex_node[0][]);
    m_problem.y = new double[vy.size()];
    for (int i = 0; i < vy.size(); i++)
      m_problem.y[i] = ((Double) vy.get(i)).doubleValue();

    // print counts
    for (int i = 0; i < counts.length; i++)
      System.out.println(counts[i] + " pairs with label " + i);

    return true;
  }


  protected boolean updateProblem(IntexDocument intexDoc, Vector vx, Vector vy)
  {
    for (int i = 0; i < intexDoc.m_sents.length; i++) {
      IntexSentence sent = intexDoc.m_sents[i];
      // read upper-triangular matrix of entity pairs
      for (int j = 0; j < sent.m_ems.length; j++) {
	ACEEMention em1 = sent.m_ems[j];
	for (int k = j + 1; k < sent.m_ems.length; k++) {
	  ACEEMention em2 = sent.m_ems[k];
	  if (sent.m_pairs[j][k] != null) {
	    String[][] vecFeat = getFeatureVector(sent, em1, em2);
	    if (vecFeat != null) {
	      // store class information
	      vy.add(new Double(sent.m_pairs[j][k].m_nTrueLabel));

	      intex_node[] inode = new intex_node[1];
	      inode[0] = new intex_node();
	      inode[0].m_value = vecFeat;
	    
	      // store feature vector
	      vx.add(inode);
	      
	      /*
	    if (sent.m_pairs[j][k].m_nTrueLabel != EntityPair.Labels.NONE)
	      System.out.print("{+++}: ");
	    else
	      System.out.print("{---}: ");

	    for (int w = 0; w < inode[0].m_value.length; w++) {
	      if (inode[0].m_value[w][FeatureDictionary.FEAT_ETYPE] != null)
		System.out.print("[");
	      System.out.print(inode[0].m_value[w][0]);
	      if (inode[0].m_value[w][FeatureDictionary.FEAT_ETYPE] != null)
		System.out.print("]");
	      if (w < inode[0].m_value.length - 1)
		System.out.print(" --- ");
	    }
	    System.out.println();
	      */
	    }
	  }
	}
      }
    }

    return true;
  }


  public String[][] getFeatureVector(IntexSentence sent, 
				     ACEEMention em1, ACEEMention em2)
  {
    int e1 = em1.m_nHBeg + em1.m_nHLen - 1;
    int e2 = em2.m_nHBeg + em2.m_nHLen - 1;

    assert e1 < e2;

    int b1 = em1.m_nBeg;
    int b2 = em2.m_nBeg;
    
    int[] path;
    // "the American bush people"
    if (b1 >= b2) {
      boolean bPos = false;
      for (int i = e1; !bPos && i < e2; i++)
	if (sent.m_aceTS.m_vecPOS[i].equals("POS"))
	  bPos = true;

      if (!bPos) {
	path = new int[2];
	path[0] = e1;
	path[1] = e2;
      }
      else
	path = sent.m_aceTS.getShortestPath(e1, e2);
    }
    else
      path = sent.m_aceTS.getShortestPath(e1, e2);

    if (path == null)
      path = sent.m_aceTS.getShortPath(e1, e2);

    // return null for long pahts 
    //    if (path.length > MAX_LENGTH)
    //      return null;

    String[][] vecRel = new String[path.length][];

    // store words and POS tags
    for (int i = 0; i < path.length; i++) {
      String[] vecFeat = new String[FeatureDictionary.ftCount()];
      vecFeat[FeatureDictionary.FEAT_WORD] = 
	FeatureDictionary.getAddWordFeature(sent.m_aceTS.m_vecTokens[path[i]],
					    sent.m_aceTS.m_vecPOS[path[i]]);
      vecFeat[FeatureDictionary.FEAT_POS] =
	FeatureDictionary.getAddFeature(FeatureDictionary.FEAT_POS,
					sent.m_aceTS.m_vecPOS[path[i]]);
      vecRel[i] = FeatureDictionary.getDictFeatures(vecFeat);
    }
    
    // store additional features for left entity mention
    vecRel[0][FeatureDictionary.FEAT_ETYPE] =  
      FeatureDictionary.getAddFeature(FeatureDictionary.FEAT_ETYPE,
				      em1.m_entity.m_strType);
    vecRel[0][FeatureDictionary.FEAT_MTYPE] =
      FeatureDictionary.getAddFeature(FeatureDictionary.FEAT_MTYPE,
				      em1.m_strType);
  
    vecRel[vecRel.length - 1][FeatureDictionary.FEAT_ETYPE] =  
      FeatureDictionary.getAddFeature(FeatureDictionary.FEAT_ETYPE,
				      em2.m_entity.m_strType);
    vecRel[vecRel.length - 1][FeatureDictionary.FEAT_MTYPE] =
      FeatureDictionary.getAddFeature(FeatureDictionary.FEAT_MTYPE,
				      em2.m_strType);

    return vecRel;
  }
  

  protected boolean setParams()
  {
    m_params = new svm_parameter();

    // default values
    m_params.svm_type = svm_parameter.C_SVC;
    m_params.kernel_type = svm_parameter.RELATION;
    m_params.degree = 3;
    m_params.gamma = 0;
    m_params.coef0 = 0;
    m_params.nu = 0.5;
    m_params.cache_size = 256;
    m_params.C = 1;
    m_params.eps = 1e-3;
    m_params.p = 0.1;
    m_params.shrinking = 1;
    m_params.probability = 0; // 1
    m_params.nr_weight = 0;
    m_params.weight_label = new int[0];
    m_params.weight = new double[0];
    
    m_params.m_bDist = false;

    return true;
  }


  public boolean testModel(String corpus_dir, String corpus_list, String out_dir)
  {
    loadDictionary();
    
    int nTrue = 0;
    int nExt = 0;
    int nExtTrue = 0;

    try {
      BufferedReader br = new BufferedReader(new FileReader(corpus_list));
      String strLine = br.readLine();
      while (strLine != null) {
        strLine = strLine.trim();
        if (strLine.length() > 0) {
          // load input document
          String strInFile = corpus_dir + "/" + strLine;
          
          System.out.print("processing [" + strInFile + "] ... ");

          ACEDocument aceDoc = new ACEDocument();
          aceDoc.load(strInFile);
	  aceDoc.removeRepetitions();
          ACETokSentDocument aceTSDoc = new ACETokSentDocument(aceDoc);
          
          // create relation extraction document
          IntexDocument intexDoc = new IntexDocument(aceTSDoc);
	  inferRelations(intexDoc);

	  int[] stats = intexDoc.getStats();
	  nTrue += stats[0];
	  nExt += stats[1];
	  nExtTrue += stats[2];
	  System.out.println(nTrue + " " + nExt + " " + nExtTrue);
	  // write output document
	  //	  String strOutFile = out_dir + "/" + strLine + ".out";
	  //	  intexDoc.writeInf(strOutFile);

	  System.out.println("done!");
        }	  
        strLine = br.readLine();
      }
    }
    catch (IOException e) {
      System.err.println(e);
      return false;
    }
    
    System.out.println("Precision = " + nExtTrue / (double) nExt);
    System.out.println("Recall = " + nExtTrue / (double) nTrue);

    return true;
  }
  
  
  protected void loadDictionary()
  {
    for (int i = 0; i < m_model.SV.length; i++) {
      intex_node inode = (intex_node) m_model.SV[i][0];
      for (int j = 0; j < inode.m_value.length; j++)
	for (int ft = 0; ft < intex_node.FT; ft++)
	  if (inode.m_value[j][ft] != null)
	    inode.m_value[j][ft] = 
	      FeatureDictionary.getAddFeature(ft, inode.m_value[j][ft]);
    }
  }


  protected boolean inferRelations(IntexDocument intexDoc)
  {
    // label probabilities
    double[] probs = new double[2];

    for (int i = 0; i < intexDoc.m_sents.length; i++) {
      IntexSentence sent = intexDoc.m_sents[i];
      // read upper-triangular matrix of entity pairs
      for (int j = 0; j < sent.m_ems.length; j++) {
	ACEEMention em1 = sent.m_ems[j];
	for (int k = j + 1; k < sent.m_ems.length; k++) {
	  ACEEMention em2 = sent.m_ems[k];
	  if (sent.m_pairs[j][k] != null) {
	    String[][] vecFeat = getFeatureVector(sent, em1, em2);
	    if (vecFeat != null) {
	      // create relation example sentence
	      intex_node[] inode = new intex_node[1];
	      inode[0] = new intex_node();
	      inode[0].m_value = vecFeat;

	      // class inference
	      sent.m_pairs[j][k].m_nInfLabel = (int) svm.svm_predict(m_model,
								     inode);
	    }
	    else
	      sent.m_pairs[j][k].m_nInfLabel = EntityPair.Labels.NONE;
	    //	    if (sent.m_pairs[j][k].m_nInfLabel != EntityPair.Labels.NONE)
	    //	      System.out.println(sent.m_pairs[j][k].m_nInfLabel);
	  }
	}
      }
    }
    
    return true;    
  }
  

  static public void main_train(String[] args)
  {
    if (args.length != 3) {
      System.out.println("Usage: java IntexModel <corpus-dir> <corpus-list> " + 
                         "<intex-model>!");
      System.exit(1);
    }
    
    IntexModel intex = new IntexModel();
    intex.trainModel(args[0], args[1]);

    try {
      svm.svm_save_model(args[2], intex.m_model);    
    }
    catch (IOException e) {
      System.err.println(e);
    }
  }
  

  static public void main_test(String[] args)
  {
    if (args.length != 4) {
      System.out.println("Usage: java IntexModel <corpus-dir> <corpus-list> " + 
                         "<svm-model> <out-dir>!");
      System.exit(1);
    }

    IntexModel intex = new IntexModel();
    try {
      intex.m_model = svm.svm_load_model(args[2]);
    }
    catch (IOException e) {
      System.err.println(e);
    }

    intex.testModel(args[0], args[1], args[3]);
  }


  static public void main(String[] args)
  {
    if (args.length == 0) {
      System.out.println("Usage: java IntexModel -[train|test] ...");
      System.exit(1);
    }
    
    String[] newArgs = new String[args.length - 1];
    for (int i = 1; i < args.length; i++)
      newArgs[i-1] = new String(args[i]);
    
    if (args[0].compareTo("-train") == 0) {
      main_train(newArgs);
      return;
    }
    if (args[0].compareTo("-test") == 0) {
      main_test(newArgs);
      return;
    }
  }


}
